
from numpy import *
from math import *
import Gnuplot

class Plotter:
  """
  Class maintaining plotting routines. Uses Gnuplot for output.
  """

#
# Properties
#
  Filename = None  # Filename for a possible raw data output 
                   # (see, e.g., the stack-graph routine)
  onScreen = True  # Whether to plot to the screen (window) or to file...
  GraphOutput=""   # ... named GraphOutput
  GraphTitle=""    # The title of the graph

##################################################################
#
# Init
#
##################################################################
  def __init__(self):
    """
    Initializes Gnuplot.
    """

    self.graph = Gnuplot.Gnuplot()

##################################################################
#
# Plot_Fermi_surface
#
##################################################################
  def Plot_Fermi_surface(self, kr1, kr2, Akw):
    """
    Plots a spectral function Akw[k1, k2] corresponding
    to a constant energy BZ slice. 'kr1', 'kr2' contain
    actual k-point coordinates.
    """

    screen = self.onScreen
    output = self.GraphOutput
    title = self.GraphTitle

    graph = self.graph

    if screen:
      graph("set term wxt enhanced")
    else:
      assert output, "  The output filename must be provided."

      graph("set term png giant enhanced size 800,600")
      graph('set output "%s"'%(output))

    if title:
      graph('set title "%s"'%(title) )

# Initialize image plot
    graph("set pm3d map interpolate 2,2")
    graph("set size square")

    graph("gamma = 0.4")

    setstr = "set palette function"
    setstr += "(1-gray)**(1/gamma), "
    setstr += "(1-gray)**(1/gamma), "
    setstr += "(1-gray)**(1/gamma) positive"

    graph(setstr)

# Lines marking the triangular BZ
    graph("unset arrow")
    graph("set arrow from -4*pi/3,0 to -2*pi/3,2*pi/sqrt(3) nohead front")
    graph("set arrow from -2*pi/3,2*pi/sqrt(3) to  2*pi/3,2*pi/sqrt(3) nohead front")
    graph("set arrow from 2*pi/3,2*pi/sqrt(3) to 4*pi/3,0  nohead front")
    graph("set arrow from 4*pi/3,0 to  2*pi/3,-2*pi/sqrt(3) nohead front")
    graph("set arrow from 2*pi/3,-2*pi/sqrt(3) to  -2*pi/3,-2*pi/sqrt(3) nohead front")
    graph("set arrow from -2*pi/3,-2*pi/sqrt(3) to -4*pi/3,0  nohead front")

#    graph("unset colorbox")
    graph("set xr [%f:%f]"%(amin(kr1), amax(kr1)))
    graph("set yr [%f:%f]"%(amin(kr2), amax(kr2)))

    nk1 = len(kr1)
    nk2 = len(kr2)

    dat = zeros((nk2,nk1,3))

    for j1 in range(nk2):
      for j2 in range(nk1):
        dat[j1,j2,:] = array([kr1[j1], kr2[j2], Akw[j1,j2]])

    d = Gnuplot.Data(dat)

    graph.splot(d)

##################################################################
#
# Plot_SF_kpoint
#
##################################################################
  def Plot_SF_kpoint(self, ens, Akw, Efermi):
    """
    Plots the energy dependent spectral function Akw[ie] for
    a single k-point.
    """

    filename = self.Filename
    screen = self.onScreen
    output = self.GraphOutput
    title = self.GraphTitle

    graph = self.graph

    if screen:
      graph("set term wxt enhanced")
    else:
      assert output, "  The output filename must be provided."

      graph("set term postscript color enhanced 'Helvetica,18' lw 1.5")
      graph('set output "%s"'%(output))

    if title:
      graph('set title "%s"'%(title) )

    graph('set xtics autofreq 0.2')

    ne = len(ens)
    dat = zeros((ne,2))

    for j1 in range(ne):
      dat[j1,:] = array([ens[j1], Akw[j1]])

    if not filename is None:
      fout = open( filename, "wt" )
      for ie, e in enumerate( ens ):
        fout.write("%15.10f   %15.10f\n"%( e, Akw[ie] ) )
      fout.close()

    d = Gnuplot.Data(dat, with_ = "l")

    graph.plot(d)


##################################################################
#
# Plot_SF_stack
#
##################################################################
  def Plot_SF_stack(self, kpts, ens, Akw, Efermi, nk, nr, lett):
    """
    Plots stack plot of the spectral function Akw corresponding to
    a scan along the k-lines defined in 'kpts'.
      * nr   -- number of high-symmetry lines
      * lett -- labels of high-symmetry k-points
    """

    screen = self.onScreen
    output = self.GraphOutput
    title = self.GraphTitle

    graph = self.graph

    if screen:
      graph("set term wxt enhanced")
    else:
      assert output, "  The output filename must be provided."

      graph("set term postscript color enhanced 'Helvetica,18' lw 1.5")
      graph('set output "%s"'%(output))

    if title:
      graph('set title "%s"'%(title) )

    ne = len(ens)

    A = zeros((ne,nk*nr))
    dat = zeros((ne,nk*nr,3))
    kpns = zeros(nk*nr)

    xticstr = "set xtics ("

    for ir in range(nr):
      i1 = ir*nk
      i0 = i1 - 1
      i2 = (ir + 1)*nk
      if ir > 0:
        kpns[i1:i2] = kpts[i1:i2] + kpns[i0]
        xticstr += ",'%s' %f"%(lett[ir+1], kpns[i2-1])
      else:
        kpns[i1:i2] = kpts[i1:i2]
        xticstr += "'%s' 0, '%s' %f"%(lett[0], lett[1], kpns[i2-1])

      A[:,i1:i2] = Akw[:,:,ir]

    xticstr += ")"
    graph(xticstr)

    graph("unset ztics")
    graph("set view 36, 75")
#    graph("set view 46, 101")
    graph("set border 15")
    graph("set zrange [0:{0}]".format( amax(A) + 0.1 ) )
    graph("set xyplane at 0.0")
    graph("set xzeroaxis")

    ds = []
    for j2 in range(nk*nr):
      for j1 in range(ne):
        dat[j1,j2,:] = array([kpns[j2], ens[j1], A[j1,j2]])

      ds.append( Gnuplot.Data( dat[:,j2,:], with_ = "l lc 7 lt 1" ) )


    d = Gnuplot.Data(dat, with_ = "l")

#    graph("set arrow from {0},{1} to {2},{3} nohead front".\
#          format(kpns[i1],0.0,kpns[i2-1],0.0))

#    graph.splot(d)
    graph.splot(*ds)

##################################################################
#
# Plot_SF
#
##################################################################
  def Plot_SF(self, kpts, ens, Akw, Efermi, nk, nr, lett):
    """
    Plots the k-resolved spectral function Akw[iw, ik, ir].
    """

    screen = self.onScreen
    output = self.GraphOutput
    title = self.GraphTitle

    graph = self.graph

    if screen:
      graph("set term wxt enhanced")
    else:
      assert output, "  The output filename must be provided."
      print "  Output file:", output

      graph("set term png giant enhanced size 800,600")
      graph('set output "%s"'%(output))

    if title:
      graph('set title "%s"'%(title) )

#
# In this implementation n = nk*nr is assumed
#
    n = nk*nr
    ne = len(ens)

# Set the absolute scale for x-axis
    maxk = fsum([kpts[(i+1)*nk-1] for i in range(nr)])

# Initialize image plot
    graph("set pm3d map interpolate 2,2")

    graph("gamma = 0.3")

    setstr = "set palette function"
    setstr += "(1-gray)**(1/gamma), "
    setstr += "(1-gray)**(1/gamma), "
    setstr += "(1-gray)**(1/gamma) positive"

    graph(setstr)

    graph("unset colorbox")

# Initialize multiplot mode
    graph("set size 1,1")
    graph("set origin 0,0")
    graph("set multiplot")

    yr1 = ens.min()
    yr2 = ens.max()

# Adjust a bit for a nicer view
    dy = yr2 - yr1
    yr1 -= dy/20
    yr2 += dy/20

#    yr1 = -5.0
    setstr = 'set yrange [{0}:{1}]'.format(yr1, yr2)
    graph(setstr)

    graph("set ytics mirror")
    graph("set mytics")

    graph('set ylabel "e - e_f, eV"')

# Origin of a current pan: ox, oy
# Size of a current pan: sx, sy
# Everything is defined in relative coordiantes
    sy = float(1)
    oy = float(0)
    ox = float(0)
    sx = float(0)
# Width of the plot area without the margins
    lm = 0.10
    rm = 0.05
    wid = 1 - lm - rm

    graph("set lmargin at screen %f"%(lm))
#
# Put all k-points into a single ascending array
#
    A = zeros((ne,nk*nr))
    kpns = zeros(nk*nr)

    xticstr = "set xtics ("

    for ir in range(nr):
      i1 = ir*nk
      i0 = i1 - 1
      i2 = (ir + 1)*nk
      if ir > 0:
        kpns[i1:i2] = kpts[i1:i2] + kpns[i0]

        xticstr += ",'%s' %f"%(lett[ir+1], kpns[i2-1])

        if ir < nr :
          arrstr = "set arrow from %f,%f to %f,%f nohead front"\
                   %(kpns[i1], yr1, kpns[i1], yr2)
          graph(arrstr)
      else:
        kpns[i1:i2] = kpts[i1:i2]
        xticstr += "'%s' 0, '%s' %f"%(lett[0], lett[1], kpns[i2-1])

      A[:,i1:i2] = Akw[:,:,ir]

    nk = nk*nr
    nr = 1

    xticstr += ")"
    graph(xticstr)

    for ir in range(nr):
        i1 = ir*nk
        i2 = (ir + 1)*nk

        kmaxi1 = kpns[i1]
        kmaxi2 = kpns[i2 - 1]

        ox += sx

        if ir == 0:
          sx = kmaxi2/maxk*wid + lm
        elif ir == nr - 1:
          sx = kmaxi2/maxk*wid + rm
        else:
          sx = kmaxi2/maxk*wid

        setstr = 'set size {0},{1}'.format(sx,sy)
        graph(setstr)

        setstr = 'set origin {0},{1}'.format(ox,oy)
        graph(setstr)

        if ir > 0:
            graph('set lmargin 0')
            graph('set format y ""')
            graph('unset ytics')
            graph('unset ylabel')

        if ir < nr - 1:
            graph('set rmargin 0')

        else:
            graph("set rmargin at screen %f"%(1 - rm))

#
# Draw spectral function
#
        dat = zeros((ne,nk,3))

        for j1 in range(ne):
          for j2 in range(nk):
            dat[j1,j2,:] = array([kpns[j2], ens[j1], A[j1,j2]])

        d = Gnuplot.Data(dat)

        graph("set arrow from {0},{1} to {2},{3} nohead front".\
              format(kpns[i1],0.0,kpns[i2-1],0.0))

        graph.splot(d)

        graph("unset arrow")

